from src.test_utils.execute_utils import transform_torch_model_mnn, execute_cpp, execute_torch, diff_cpp_python
test_data = {
        "dataset_type": "tabular",
        "training_data_location": "./data/tabular",
        "torch_mnn_transformer_path": "./src/python/torch_mnn_transformer.py",
        "torch_trainer": "./src/python/torch_trainer.py",
        "cpp_executable_path": "./src/cpp/main_mnn_train.out",
}

if __name__ == "__main__":
    mnn_model_export_dir = "./mnn_model/"
    # 0. compile cpp program, and remove to the current directory

    # 1. transform torch model to mnn model
    transform_torch_model_mnn(test_data["torch_mnn_transformer_path"], mnn_model_export_dir, test_data["dataset_type"])
    print("Finish transform torch model to mnn model")

    # 2. execute cpp program
    mnn_4_cpp_path = mnn_model_export_dir + test_data["dataset_type"] + ".mnn"    # This file is generated by step 1!
    execute_cpp(test_data["cpp_executable_path"], test_data["dataset_type"], mnn_4_cpp_path, test_data["training_data_location"])
    print("Finish execute cpp program")

    # 3. execute python program
    execute_torch(test_data["torch_trainer"], test_data["training_data_location"], mnn_model_export_dir + test_data["dataset_type"] + ".pth", test_data["dataset_type"])
    print("Finish execute python program")

    # 4. compare the result
    diff_cpp_python(test_data["dataset_type"])
    print("Finish compare the result using diff")